# import necessary libraries
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import torchvision
import torch
import numpy as np
import cv2
import time
import glob, random, os
import matplotlib.patches as patches

# Class labels from official PyTorch documentation for the pretrained model
# Note that there are some N/A's
# for complete list check https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
# we will use the same list for this notebook
COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]


def get_prediction(img_path, threshold):
    """
    get_prediction
      parameters:
        - img_path - path of the input image
        - threshold - threshold value for prediction score
      method:
        - Image is obtained from the image path
        - the image is converted to image tensor using PyTorch's Transforms
        - image is passed through the model to get the predictions
        - class, box coordinates are obtained, but only prediction score > threshold
          are chosen.

    """
    img = Image.open(img_path)
    transform = T.Compose([T.ToTensor()])
    img = transform(img)
    pred = model([img])
    pred_class = np.array([COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())])
    pred_boxes = np.array([[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].cpu().numpy())])
    pred_score = pred[0]['scores'].cpu().numpy()
    pred_t = pred_score > threshold
    pred_boxes = pred_boxes[pred_t]
    pred_class = pred_class[pred_t]
    if len(pred_boxes) > 0:
        keypoints = pred[0]["keypoints"].cpu().squeeze(dim=0).numpy()  # for test batch_size = 1
        keypoints_score = pred[0]["keypoints_scores"].cpu().squeeze(dim=0).numpy()  # for test batch_size = 1
        return pred_boxes, pred_class, keypoints, keypoints_score
    else:
        return [],[],[],[]


def object_detection_api(img_path, threshold=0.5, rect_th=3, text_size=3, text_th=3):
    """
    object_detection_api
      parameters:
        - img_path - path of the input image
        - threshold - threshold value for prediction score
        - rect_th - thickness of bounding box
        - text_size - size of the class label text
        - text_th - thichness of the text
      method:
        - prediction is obtained from get_prediction method
        - for each prediction, bounding box is drawn and text is written
          with opencv
        - the final image is displayed
    """
    boxes, pred_cls, keypoint, keypoints_score = get_prediction(img_path, threshold)
    if len(boxes) > 0:
        img = np.asarray(Image.open(img_path).convert("RGB"))
        fig, ax = plt.subplots(1)
        ax.imshow(img)
        for i in range(len(boxes)):
            xmin, ymin = boxes[i][0]
            xmax, ymax = boxes[i][1]
            rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=2, edgecolor='r',
                                     facecolor='none')
            ax.add_patch(rect)
            ax.scatter(keypoint[i][:, 0], keypoint[i][:, 1], c='green', marker='s')
        plt.show()




def check_inference_time(image_path, gpu=False):
    img = Image.open(image_path)
    transform = T.Compose([T.ToTensor()])
    img = transform(img)
    if gpu:
        model.cuda()
        img = img.cuda()
    else:
        model.cpu()
        img = img.cpu()
    start_time = time.time()
    pred = model([img])
    end_time = time.time()
    return end_time - start_time


if __name__ == "__main__":
    # get the pretrained model from torchvision.models
    # Note: pretrained=True will get the pretrained weights for the model.
    # model.eval() to use the model for inference
    from pycocotools.coco import COCO
    coco = COCO("/media/retoo/RetooDisk/wanghui/Data/coco2017/annotations/person_keypoints_val2017.json")
    img_id = coco.imgs.keys()
    model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)
    model.eval()
    with torch.no_grad():
        for id in list(img_id)[1:]:
            ann_ids = coco.getAnnIds(imgIds=id)
            target = coco.loadAnns(ann_ids)

            path = os.path.join("/media/retoo/RetooDisk/wanghui/Data/coco2017/val2017",coco.loadImgs(id)[0]['file_name'])

            object_detection_api(path, rect_th=4, text_th=2, text_size=1, threshold=0.6)
            #
            # cpu_time = sum([check_inference_time('./girl_cars.jpg', gpu=False) for _ in range(10)]) / 10.0
            # gpu_time = sum([check_inference_time('./girl_cars.jpg', gpu=True) for _ in range(10)]) / 10.0
            #
            # print('\n\nAverage Time take by the model with GPU = {}s\nAverage Time take by the model with CPU = {}s'.
            #       format(gpu_time, cpu_time))
